import warnings
import numpy as np
from matplotlib import pyplot as plt
from scipy.optimize import fmin, fminbound, brute
from scipy.interpolate import splrep
from scipy.interpolate import splev
from simulate import ResSim
from csv import reader
from datetime import date
from pdb import set_trace
"""
Module to import, analyze, and plot data obtained 
from impact measurements
"""

# Global constants
g = 9.8             # gravity (m/s^2)

class Measurement():
    """
    Class to hold an impact measurement data and provide methods for analysis.
    Class variables:
        time - numpy array of measurement time axis 
        voltage - numpy array of measurement voltage
        params - dictionary experimental parameters 
    """

    def __init__(self):
        """ 
        Set default values for the class variables
        """
        self.name = None
        self.ExpDescription = None
        self.SavePath = None
        self.raw_data = np.zeros((2,1000)) # [time (s), voltage (mV)]
        self.params = {
                'h' : None,            # drop height (m)
                'R' : None,            # sphere radius (m)
                'Es' : None,           # sphere modulus (Pa)
                'vs' : None,           # sphere Poisson's ratio
                'm' : None,            # sphere mass (kg)
                'cal' : None,          # calibration constant
                'm_crystal' : None,    # mass of the crystal/compact/load cell (g)
                'k_crystal' : None,    # stiffness of the spring
                'mu_dashpot' : None    # damping factor of the crystal
                 }
        self.estimates = {
                'Ep' : None,           # plane modulus (Pa)
                'vp' : None,           # plane Poisson's ratio
                'COR' : None,          # coefficient of restitution
                'R2' : None            # Sum of squares of residuals (R2)
                 }
        self.other_data = {}            # place holder for user added data
        self.figure = {                       # dictionary to keep track of plot objects
                'fig' : None,                 # matplotlib figure object (figure of the data)
                'ax' : None,                  # matplotlib Axes object
                'line_raw' : None,            # matplotlib Line2D object
                'line_fit' : None             # matplotlib Line2D object
                }

    @classmethod
    def from_file(cls, fname):
        """
        Load a previously saved measurement file
        """
        meas = cls()

        gui_vals = {}
        data = []
        lineNo = {'Begin Parameter':0,'End Parameter':0,
                  'Begin Data':0, 'End Data':0}

        # Determine positions of data in file
        with open(fname, mode='r') as f:
            for ln, txt in enumerate(f):
                if txt.find("****Begin Parameter Values****") != -1:
                    lineNo['Begin Parameter'] = ln
                elif txt.find("****End Parameter Values****") != -1: 
                    lineNo['End Parameter'] = ln
                elif txt.find("****Begin Data****") != -1: 
                    lineNo['Begin Data'] = ln + 1
                elif txt.find("****End Data****") != -1: 
                    lineNo['End Data'] = ln
                
            # Read in the file data using the determined positions
            f.seek(0)
            csv_f = reader(f)
            for ln, txt in enumerate(csv_f):
                if (ln > lineNo['Begin Parameter'] and 
                    ln < lineNo['End Parameter']):
                    k, v = txt[0].split("=:=")
                    try:
                        k, v = k.strip(), float(v.strip())
                    except ValueError:
                        k, v = k.strip(), str(v.strip())
                    gui_vals.update({k : v})
                elif (ln > lineNo['Begin Data'] and 
                      ln < lineNo['End Data']):
                    data.append(txt)

        # Create a dictionary of the parameter values
        exp_params = {'h':gui_vals['DropHeight'] * 1.0e-2,
                  'm':gui_vals['SphereMass'] * 1.0e-3,
                  'R':gui_vals['SphereRadius'] * 1.0e-3,
                  'Es':gui_vals['SphereModulus'] * 1.0e9,
                  'vs':gui_vals['SpherePR'],
                  'cal' : gui_vals['Calibration'] * 1.0e-9,
                  'm_crystal':gui_vals['CrystalMass'] * 1e-3,
                  'k_crystal':gui_vals['CrystalStiffness'] * 1e9,
                  'mu_dashpot':gui_vals['CrystalDamping'] * 1e3}
        est_params = {'Ep':gui_vals['SurfaceModulus'] * 1.0e9,
                  'vp':gui_vals['SurfacePR'],
                  'COR':gui_vals['COR']}

        # Save the experimental description and path
        meas.name = fname.split('/')[-1]
        meas.ExpDescription = gui_vals['ExpDescription']
        meas.SavePath = gui_vals['SavePath']

        # Create array of the time series voltage
        data = np.array(data,dtype='f8').transpose()

        # Set the data and parameters
        meas.set_data(data[0] * 1.0e-6, data[1])
        meas.set_parameters(exp_params)
        meas.set_parameters(est_params)
        return meas

    @classmethod
    def from_hdf5(cls, f_h5py, hdf5_path):
        """
        Load measurement data from a hdf5 file, and create a new
        measurement instance from the data.
        f_h5py - f_h5py file object
        hdf5_path - root path that contains the measurement
                    data
        """
        if hdf5_path[0] != '/':
            hdf5_path = '/' + hdf5_path
        if hdf5_path[-1] != '/':
            hdf5_path = hdf5_path + '/'

        meas = cls()
        for atts in f_h5py[hdf5_path].attrs.items():
            if atts[0].find('Experimental Description') == 0:
                meas.ExpDescription = atts[1]
            elif atts[0].find('Original Save Path') == 0:
                meas.SavePath = atts[1]

        meas.name = f_h5py[hdf5_path+'name'].value

        meas.raw_data[0] = np.array(f_h5py[hdf5_path+'time'].value)
        meas.raw_data[1] = np.array(f_h5py[hdf5_path+'voltage'].value)

        params = {}
        for k, v in f_h5py[hdf5_path].items():
            if k == "time" or k == "voltage":
                pass
            else:
                params.update({str(k) : v.value})
        meas.set_parameters(params)
        return meas

    @classmethod
    def from_array(cls, marray):
        m = cls()
        m.name = marray['name'].item()
        m.raw_data[0] = marray['time']
        m.raw_data[1] = marray['voltage']
        params = {k:marray['parameters'][k].item() for k in 
                  marray['parameters'].dtype.names}
        m.set_parameters(params)
        return m

    def to_hdf5(self, f_h5py, hdf5_path):
        """
        Save the experiment to an hdf5 file
        f_h5py - h5py object file
        hdf5_path - path to create in the hdf5 file to save data
        """
        name = self.name

        if hdf5_path[0] != '/':
            hdf5_path = '/' + hdf5_path
        if hdf5_path[-1] != '/':
            hdf5_path = hdf5_path + '/'
        if name[0] == '/':
            name = name[1:]
        if name[-1] != '/':
            name = name + '/'

        full_name = hdf5_path + name
        f_h5py.create_group(full_name)
        f_h5py[full_name].attrs['Experimental Description'] = self.ExpDescription
        f_h5py[full_name].attrs['Original Save Path'] = self.SavePath
        f_h5py[full_name].attrs['Date Modified'] = date.today().__str__()

        f_h5py.create_dataset(full_name+'name', data=self.name)

        f_h5py.create_dataset(full_name+'time', data=self.raw_data[0])
        f_h5py.create_dataset(full_name+'voltage', data=self.raw_data[1])

        for k, v in dict(self.params.items() + 
                         self.estimates.items() + 
                         self.other_data.items()).iteritems():
            if v is None:
                v = "None"
            f_h5py.create_dataset(full_name+k, data=v)

    def to_array(self):
        """
        Create a numpy array of the measurement voltage
        """
        pdtype = []
        for k, v in self.get_parameters().iteritems():
            if isinstance(v, str):
                pdtype.append((k,'|S30'))
            elif isinstance(v, int):
                pdtype.append((k,'int32'))
            elif isinstance(v, float):
                pdtype.append((k,'float64'))
            else:
                pdtype.append((k,'|S30'))
        dtype = [('name','|S50'), ('parameters',pdtype), ('time','<f4',(1000,)), 
                 ('voltage','<f4',(1000,))]
        starray = np.zeros((1), dtype=dtype)
        starray['name'] = self.name
        starray['time'] = self.raw_data[0]
        starray['voltage'] = self.raw_data[1]
        for k, v in self.get_parameters().iteritems():
            starray['parameters'][k] = v
        return starray

    def set_data(self, time_axis, voltage):
        """
        Set the time axis and measured voltage
        time_axis - numpy array time (s)
        voltage - numpy array measured transient impact voltage (mV)
        """
        self.raw_data[0] = time_axis
        self.raw_data[1] = voltage

    def set_parameters(self, parameters):
        """
        Set the measurement parameters
        parameters - dictionary of parameters
        """
        for k, v in parameters.iteritems():
            if self.params.has_key(k):
                self.params.update({k:v})
            elif self.estimates.has_key(k):
                self.estimates.update({k:v})
            elif self.other_data.has_key(k):
                self.other_data.update({k:v})
            else:
                self.other_data.update({k:v})
                warnings.warn('Previously undefined parameters %s'%k)

    def add_data(self, data):
        """
        Add additional data to measurement dictionary, other_data
        data - dictionary of additional data / parameters
        """
        for k, v in data.iteritems():
            self.other_data.update({k:v})

    def get_data(self):
        """
        Return the voltage data
        """
        return self.raw_data

    def get_parameters(self):
        """
        Return the dictionary of all the parameters
        {params + estimates + other_data}
        """
        params = dict(self.params.items() + self.estimates.items() 
                      + self.other_data.items())
        return params

    def get_estimates(self):
        """
        Return the class dictionary of estimates
        """
        return self.estimates

    def crop_data(self, t, signal):
        """
        Crop the transient signal data to the portion
        where the sphere is in contact with the surface.
        t - numpy array time (s)
        signal - numpy array voltage (mV)
        return [nx1] numpy array of cropped signal
        """
        t_max = signal.argmax()
        v_max = np.max(signal)
        t_start = t_max - 5
        while np.average(signal[t_start:t_start+3]) > np.average(signal[t_start-3:t_start]) \
                and np.average(signal[t_start:t_start+3]) > 0.01*signal[t_max] \
                and t_start > 0:
            t_start = t_start - 1
        t_end = t_max + (np.sum(signal>v_max/2))/2
        while np.average(signal[t_end-3:t_end]) > np.average(signal[t_end:t_end+3]) \
                and np.average(signal[t_end-3:t_end]) > 0 \
                and t_end < t.shape[0]:
            t_end = t_end + 1
        return [t[t_start:t_end], signal[t_start:t_end]]
        
    def calc_cor(self):
        """
        Calculate the coefficient of restitution from the force curve.
        t - [nx1] numpy array of time (s)
        f - [nx1] numpy array of transient force (Pa)
        """
        sim = ResSim()
        params = dict(self.params.items() + self.estimates.items())
        sim.set_parameters({i:params[i] for i in ['h','R','Es','Ep','vs',
            'vp','m', 'COR','m_crystal','k_crystal','mu_dashpot']})
        [t, y, f] = sim.solve()

        m = self.params['m']
        h = self.params['h']

        integral = np.trapz(f,t)
        COR = 1./(m*np.sqrt(2.*g*h)) * integral - 1.
        return COR

    def fit_dynamic_force(self,time_axis,force,E1,E2,v1,v2,m,R,h,COR,plot=False):
        """
        Fit the dynamic load cell force with the theory to extract the modulus,
        E1, and Poisson ratio, v1.
        time_axis - time (sec)
        force - dynamic force curve (N)
        h - scalar drop height (m)
        R - scalar tablet Radius (m)
        E1 - scalar tablet modulus guess (Pa)
        E2 - scalar plane modulus (Pa)
        v1 - scalar tablet Poisson ratio guess
        v2 - scalar plane Poisson ratio
        m - scalar tablet mass (kg)
        COR - scalar coefficient of restitution (Not guess)
        plot - boolean True if want a plot of fit
        return fit results
        """
        def objective(p,time_axis,force,E2,v2,m,R,h):
            [E1,v1] = p
            td_simulate = ImpactSimulation()
            td_simulate.set_parameters(h,R,E1,E2,v1,v2,m,COR)
            [t,y,f] = td_simulate.solve()
            t = t - t[np.argmax(f)]
            t = t[f>0.2*np.max(f)]
            f = f[f>0.2*np.max(f)]
            spline = splrep(t,f)
            f = splev(time_axis,spline)
            residuals = np.sum((force - f)**2)
            return residuals
        def objective_old(p,time_axis,force,E2,v2,m,R,h):
            [E1,v1] = p
            td_simulate = ImpactSimulation()
            td_simulate.set_parameters(h,R,E1,E2,v1,v2,m,COR)
            [t,y,f] = td_simulate.solve()
            t = t[f>2]
            f = f[f>2]
            t = t-t[0]
            f = f-f[0]
            spline = splrep(t,f)
            f = splev(time_axis,spline)
            residuals = np.sum((force - f)**2)
            return residuals
        def plot_results(time_axis,force,E1,E2,v1,v2,m,R,h):
            td_simulate = ImpactSimulation()
            td_simulate.set_parameters(h,R,E1,E2,v1,v2,m,COR)
            [t,y,f] = td_simulate.solve()
            t = t - t[np.argmax(f)]
            t = t[f>0.2*np.max(f)]
            f = f[f>0.2*np.max(f)]
            spline = splrep(t,f)
            f = splev(time_axis,spline)
            plt.figure()
            plt.plot(time_axis,force)
            plt.plot(time_axis,f)
        #time_axis = time_axis[force>2]
        #force = force[force>2]
        #time_axis = time_axis - time_axis[0]
        #force = force - force[0]
        time_axis = time_axis - time_axis[np.argmax(force)]
        time_axis = time_axis[force>0.2*np.max(force)]
        force = force[force>0.2*np.max(force)]
        p_guess = [E1,v1]
        func_args = (time_axis,force,E2,v2,m,R,h)
        fit_results = fmin(objective,p_guess,args=func_args,xtol=10,ftol=10,maxiter=200,maxfun=None,full_output=0,disp=1,retall=1)
        if plot == True:
            E1 = fit_results[0][0]
            v1 = fit_results[0][1]
            plot_results(time_axis,force,E1,E2,v1,v2,m,R,h)
        return fit_results

    def fit_gaussian(self, t, f, plot=False):
        """
        function to fit a transint impact force signal with
        as Gaussian function
        t - [nx1] numpy array of time
        f - [nx1] numpy array of force
        plot - boolean to create plot of fit
        return list [fit parameter A, fit parameter sigma2, fit of force]
        """
        t_trans = t - t[np.argmax(f)]
        t_trans = t_trans[f>0.3*np.max(f)]
        f_trans = f[f>0.3*np.max(f)]
        fit = np.polyfit(t_trans**2,np.log(f_trans),1)
        sigma2 = -1/fit[0]
        A = np.exp(fit[1])
        f_fit = A * np.exp(-1*(t-t[np.argmax(f)])**2/sigma2)
        if plot==True:
            fig = plt.figure()
            ax = fig.add_subplot(111)
            ax.plot(t,f,color='black')
            ax.plot(t,f_fit,color='black',ls=':')
            ax.set_xlabel(r'time ($\mu$s)')
            ax.set_ylabel(r'force (N)')
            ax.legend(['experimental','Gaussian fit'])
        return [A, sigma2, f_fit]

    def fit_voltage(self):
        """
        Fit the dynamic load cell voltage signal with the Hertzian and deflectable 
        crystal equations to extract the compact modulus, Ep, and Poisson ratio, vp.
        """
        sim = ResSim()
        t_exp, v_exp = self.crop_data(self.raw_data[0].copy(), 
                       self.raw_data[1].copy())
        params = dict(self.params.items() + self.estimates.items())
        d_exp = v_exp * params['cal']
        t_exp -= t_exp[np.argmax(d_exp)]

        def objective(p,t_exp,d_exp,params):
            params.update({'Ep':p[0], 'vp':p[1], 'COR':p[2]})
            sim.set_parameters({i:params[i] for i in ['h','R','Es','Ep','vs',
                'vp','m', 'COR','m_crystal','k_crystal','mu_dashpot']})
            [t,y,f] = sim.solve()
            t_sim, d_sim = self.crop_data(t,y[:,2])
            t_sim -= t_sim[np.argmax(d_sim)]
            spline = splrep(t_sim,d_sim)
            d_sim = splev(t_exp,spline)
            residuals = np.sum((d_exp - d_sim)**2)
            return residuals

        func_args = (t_exp, d_exp, params) 
        p_guess = np.array([self.estimates['Ep'], self.estimates['vp'], 
                            self.estimates['COR']])
        p_opt = fmin(objective, p_guess, args=func_args, xtol=10, ftol=10, 
                     maxiter=200, maxfun=50, full_output=1, disp=0, retall=0)

        self.estimates.update({'Ep':p_opt[0][0], 'vp':p_opt[0][1], 
                               'COR':p_opt[0][2]})

        SSres = p_opt[1]
        SStot = np.sum((d_exp - np.mean(d_exp))**2)
        R2 = 1 - SSres / SStot
        self.estimates.update({'R2' : R2})

    def fit_piezo_parameters(self):
        """
        Fit the voltage by varying the load cell parameters
        """
        sim = ResSim()
        t_exp, v_exp = self.crop_data(self.raw_data[0].copy(), 
                       self.raw_data[1].copy())
        params = dict(self.params.items() + self.estimates.items())
        d_exp = v_exp * params['cal']
        t_exp -= t_exp[np.argmax(d_exp)]

        def objective(p,t_exp,d_exp,params):
            params.update({'m_crystal':p[0], 'k_crystal':p[1], 'mu_dashpot':p[2]})
            sim.set_parameters({i:params[i] for i in ['h','R','Es','Ep','vs',
                'vp','m', 'COR','m_crystal','k_crystal','mu_dashpot']})
            [t,y,f] = sim.solve()
            t_sim, d_sim = self.crop_data(t,y[:,2])
            t_sim -= t_sim[np.argmax(d_sim)]
            spline = splrep(t_sim,d_sim)
            d_sim = splev(t_exp,spline)
            residuals = np.sum((d_exp - d_sim)**2)
            return residuals

        func_args = (t_exp, d_exp, params) 
        p_guess = np.array([self.params['m_crystal'], self.params['k_crystal'], 
                            self.params['mu_dashpot']])
        p_opt = fmin(objective, p_guess, args=func_args, xtol=10, ftol=10, 
                     maxiter=200, maxfun=50, full_output=1, disp=0, retall=0)

        self.params.update({'m_crystal':p_opt[0][0], 'k_crystal':p_opt[0][1], 
                               'mu_dashpot':p_opt[0][2]})

        SSres = p_opt[1]
        SStot = np.sum((d_exp - np.mean(d_exp))**2)
        R2 = 1 - SSres / SStot
        self.estimates.update({'R2' : R2})

    def plot_data(self):
        """
        Plot the experimental
        """
        self.figure['fig'] = plt.figure()
        self.figure['ax'] = self.figure['fig'].add_subplot(111)
        self.figure['line_raw'] = self.figure['ax'].plot(self.raw_data[0]*1e6, 
                self.raw_data[1], color='blue', ls='-')[0]
        self.figure['ax'].set_xlabel(r'time ($\mu$s)')
        self.figure['ax'].set_ylabel(r'signal (mV)')
        plt.draw()
        plt.show()

    def plot_est(self):
        """
        Plot the simulated signal based estimated parameters
        """
        [t, y, f] = self.simulate()

        t_s = t
        v_s = y[:,2] / self.params['cal']

        t, v = self.raw_data.copy()
        t_s = t_s + t[v.argmax()] - t_s[v_s.argmax()]
        t_s *= 1.e6

        if self.figure['fig'] == None:
            self.plot_data()

        if self.figure['line_fit'] is None:
            self.figure['line_fit'] = self.figure['ax'].plot(t_s, v_s, 
                    color='green', ls='--')[0]
        else:
            self.figure['line_fit'].set_data(t_s, v_s)
        plt.draw()

    def simulate(self):
        """
        Calculate the theoretical curve based on the measurement 
        parameters and the estimated parameters
        """
        sim = ResSim()

        params = dict(self.params.items() + self.estimates.items())
        sim.set_parameters({i:params[i] for i in ['h','R','Es','Ep','vs',
            'vp','m', 'COR','m_crystal','k_crystal','mu_dashpot']})
        [t, y, f] = sim.solve()

        return [t, y, f]
